package com.xzxy.xmlg_backend.controller;

import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.tags.Tag;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatResponse;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.ai.deepseek.DeepSeekChatModel;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.http.MediaType;
import org.springframework.web.bind.annotation.*;
import org.springframework.web.servlet.mvc.method.annotation.SseEmitter;
import jakarta.servlet.http.HttpServletResponse;
import reactor.core.publisher.Flux;

import java.util.Map;
import java.util.HashMap;
import java.util.concurrent.CompletableFuture;
import java.io.IOException;

/**
 * @Author: zys
 * @CreateTime: 2025-06-17
 */


@RestController
@RequestMapping("/ai")
@Tag(name = "ai管理接口")
public class ChatController {


    private final DeepSeekChatModel chatModel;


    @Autowired
    public ChatController(DeepSeekChatModel chatModel) {
        this.chatModel = chatModel;
    }

    @Operation(summary = "直接输出",description = "块输出")
    @GetMapping("/generate")
    public Map generate(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        return Map.of("generation", chatModel.call(message));
    }

    @Operation(summary = "流式输出",description = "流式输出")
    @GetMapping("/generateStream")
    public Flux<ChatResponse> generateStream(@RequestParam(value = "message", defaultValue = "Tell me a joke") String message) {
        var prompt = new Prompt(new UserMessage(message));
        return chatModel.stream(prompt);
    }

    @Operation(summary = "推理过程流式输出", description = "返回DeepSeek推理模型的推理过程和最终答案")
    @GetMapping(value = "/generateStreamWithReasoning", produces = MediaType.TEXT_EVENT_STREAM_VALUE)
    public SseEmitter generateStreamWithReasoning(
            @RequestParam(value = "message", defaultValue = "Tell me a joke") String message,
            @RequestHeader(value = "token", required = false) String token,
            HttpServletResponse response) {
        
        // 设置SSE响应头
        response.setContentType("text/event-stream");
        response.setCharacterEncoding("UTF-8");
        response.setHeader("Cache-Control", "no-cache");
        response.setHeader("Connection", "keep-alive");
        
        SseEmitter emitter = new SseEmitter(60000L); // 60秒超时
        
        CompletableFuture.runAsync(() -> {
            try {
                var prompt = new Prompt(new UserMessage(message));
                Flux<ChatResponse> responseFlux = chatModel.stream(prompt);
                
                StringBuilder reasoningContent = new StringBuilder();
                StringBuilder finalContent = new StringBuilder();
                
                responseFlux.subscribe(
                    response1 -> {
                        try {
                            if (response1.getResults() != null && !response1.getResults().isEmpty()) {
                                var result = response1.getResults().get(0);
                                var output = result.getOutput();
                                
                                // 检查是否有推理内容 - DeepSeek推理模型会在metadata中返回reasoning_content
                                if (output.getMetadata() != null) {
                                    Object reasoningObj = output.getMetadata().get("reasoning_content");
                                    if (reasoningObj != null) {
                                        String reasoning = reasoningObj.toString();
                                        if (!reasoning.isEmpty()) {
                                            reasoningContent.append(reasoning);
                                            
                                            // 发送累积的推理过程
                                            Map<String, Object> reasoningData = new HashMap<>();
                                            reasoningData.put("type", "reasoning");
                                            reasoningData.put("content", reasoningContent.toString());
                                            reasoningData.put("isComplete", false);
                                            
                                            emitter.send(SseEmitter.event()
                                                .name("reasoning")
                                                .data(reasoningData));
                                        }
                                    }
                                    
                                    // 也检查其他可能的推理字段
                                    Object thinkingObj = output.getMetadata().get("thinking");
                                    if (thinkingObj != null) {
                                        String thinking = thinkingObj.toString();
                                        if (!thinking.isEmpty()) {
                                            reasoningContent.append(thinking);
                                            
                                            // 发送累积的推理过程
                                            Map<String, Object> reasoningData = new HashMap<>();
                                            reasoningData.put("type", "reasoning");
                                            reasoningData.put("content", reasoningContent.toString());
                                            reasoningData.put("isComplete", false);
                                            
                                            emitter.send(SseEmitter.event()
                                                .name("reasoning")
                                                .data(reasoningData));
                                        }
                                    }
                                }
                                
                                // 检查最终内容
                                String content = output.getText();
                                if (content != null && !content.isEmpty()) {
                                    finalContent.append(content);
                                    
                                    // 发送累积的最终答案内容
                                    Map<String, Object> contentData = new HashMap<>();
                                    contentData.put("type", "content");
                                    contentData.put("content", finalContent.toString());
                                    contentData.put("isComplete", false);
                                    
                                    emitter.send(SseEmitter.event()
                                        .name("content")
                                        .data(contentData));
                                }
                            }
                        } catch (IOException e) {
                            emitter.completeWithError(e);
                        }
                    },
                    error -> {
                        emitter.completeWithError(error);
                    },
                    () -> {
                        try {
                            // 发送完成信号
                            Map<String, Object> completeData = new HashMap<>();
                            completeData.put("type", "complete");
                            completeData.put("reasoningContent", reasoningContent.toString());
                            completeData.put("finalContent", finalContent.toString());
                            completeData.put("isComplete", true);
                            
                            emitter.send(SseEmitter.event()
                                .name("complete")
                                .data(completeData));
                            
                            emitter.complete();
                        } catch (IOException e) {
                            emitter.completeWithError(e);
                        }
                    }
                );
                
            } catch (Exception e) {
                emitter.completeWithError(e);
            }
        });
        
        return emitter;
    }
}